import numpy as np
from typing import List, Tuple
from ..mathq.useful_math import dyn_pseudoInv, pseudoInv_right_weighted

class Task:
    def __init__(self, name: str):
        self.taskName: str = name
        self.id: int = 0
        self.parentId: int = 0
        self.childId: int = 0
        self.dxDes = np.empty(0, dtype=np.float64)
        self.ddxDes = np.empty(0, dtype=np.float64)
        self.delta_q = np.empty(0, dtype=np.float64)
        self.dq = np.empty(0, dtype=np.float64)
        self.ddq = np.empty(0, dtype=np.float64)
        self.J = np.empty((0, 0), dtype=np.float64)
        self.dJ = np.empty((0, 0), dtype=np.float64)
        self.Jpre = np.empty((0, 0), dtype=np.float64)
        self.N = np.empty((0, 0), dtype=np.float64)
        self.kp = np.empty((0, 0), dtype=np.float64)
        self.kd = np.empty((0, 0), dtype=np.float64)
        # self.W = np.diag([1, 2, 3])  # weighted matrix for pseudo inverse
        self.W = np.empty((0, 0), dtype=np.float64)
        self.errX = np.empty(0, dtype=np.float64)
        self.derrX = np.empty(0, dtype=np.float64)

class PriorityTasks:
    def __init__(self):
        self.taskLib: List[Task] = []
        self.nameList: List[str] = []
        self.idList: List[int] = []
        self.parentIdList: List[int] = []
        self.childIdList: List[int] = []
        self.out_delta_q = np.empty((0, 0), dtype=np.float64)
        self.out_dq = np.empty((0, 0), dtype=np.float64)
        self.out_ddq = np.empty((0, 0), dtype=np.float64)
        self.startId: int = 0

    def addTask(self, name: str):
        self.taskLib.append(Task(name))
        self.taskLib[-1].id = len(self.taskLib) - 1
        self.nameList.append(name)

    def getId(self, name: str) -> int:
        try:
            return self.nameList.index(name)
        except ValueError:
            print(f"Cannot find wbc task: {name}!")
            return -1

    def buildPriority(self, taskOrder: List[str]):
        self.startId = self.getId(taskOrder[0])
        for i, task in enumerate(taskOrder):
            idCur = self.getId(task)
            if i == 0:
                self.taskLib[idCur].parentId = -1
            else:
                idBef = self.getId(taskOrder[i-1])
                self.taskLib[idCur].parentId = idBef
            if i == len(taskOrder) - 1:
                self.taskLib[idCur].childId = -1
            else:
                idNxt = self.getId(taskOrder[i+1])
                self.taskLib[idCur].childId = idNxt

    def printTaskInfo(self):
        for task in self.taskLib:
            print("-------------")
            print(f"taskName={task.taskName}")
            print(f"parentId={task.parentId}, childId={task.childId}")

    def computeAll(self, des_delta_q: np.ndarray, des_dq: np.ndarray, des_ddq: np.ndarray, 
                   dyn_M: np.ndarray, dyn_M_inv: np.ndarray, dq: np.ndarray):
        curId = self.startId
        parentId = self.taskLib[curId].parentId
        childId = self.taskLib[curId].childId

        for _ in range(len(self.taskLib)):
            if parentId == -1:
                self.taskLib[curId].N = np.eye(self.taskLib[curId].J.shape[1])
                self.taskLib[curId].Jpre = self.taskLib[curId].J @ self.taskLib[curId].N
                self.taskLib[curId].delta_q = des_delta_q + pseudoInv_right_weighted(self.taskLib[curId].Jpre, self.taskLib[curId].W) @ self.taskLib[curId].errX
                self.taskLib[curId].dq = des_dq
                ddxcmd = self.taskLib[curId].ddxDes + self.taskLib[curId].kp @ self.taskLib[curId].errX + self.taskLib[curId].kd @ self.taskLib[curId].derrX
                self.taskLib[curId].ddq = des_ddq + dyn_pseudoInv(self.taskLib[curId].Jpre, dyn_M_inv, True) @ (ddxcmd - self.taskLib[curId].J @ dq)
            else:
                self.taskLib[curId].N = self.taskLib[parentId].N @ (np.eye(self.taskLib[parentId].Jpre.shape[1]) - 
                                        pseudoInv_right_weighted(self.taskLib[parentId].Jpre, self.taskLib[parentId].W) @ self.taskLib[parentId].Jpre)
                self.taskLib[curId].Jpre = self.taskLib[curId].J @ self.taskLib[curId].N
                self.taskLib[curId].delta_q = self.taskLib[parentId].delta_q + pseudoInv_right_weighted(self.taskLib[curId].Jpre, self.taskLib[curId].W) @ (self.taskLib[curId].errX - self.taskLib[curId].J @ self.taskLib[parentId].delta_q)
                self.taskLib[curId].dq = self.taskLib[parentId].dq + pseudoInv_right_weighted(self.taskLib[curId].Jpre, self.taskLib[curId].W) @ (self.taskLib[curId].dxDes - self.taskLib[curId].J @ self.taskLib[parentId].dq)
                ddxcmd = self.taskLib[curId].ddxDes + self.taskLib[curId].kp @ self.taskLib[curId].errX + self.taskLib[curId].kd @ self.taskLib[curId].derrX
                self.taskLib[curId].ddq = self.taskLib[parentId].ddq + dyn_pseudoInv(self.taskLib[curId].Jpre, dyn_M_inv, True) @ (ddxcmd - self.taskLib[curId].dJ @ dq - self.taskLib[curId].J @ self.taskLib[parentId].ddq)

            if childId != -1:
                parentId = curId
                curId = childId
                childId = self.taskLib[curId].childId
            else:
                break

        self.out_delta_q = self.taskLib[curId].delta_q
        self.out_dq = self.taskLib[curId].dq
        self.out_ddq = self.taskLib[curId].ddq